#!/usr/bin/python2.5
# Copyright (c) 2008 MalFind

__author__ = "noe.lutz@gmail.com (Noe Lutz)"

import re
import logging
import urlparse
import urllib

class LookupUrlGenerator:
  """ Given a url, generates a list of lookup urls for the Google safe browsing
      API.  For more imformation please see:
      http://code.google.com/p/google-safe-browsing/wiki/Protocolv2Spec
      """
  def Generate(self, url):
    """ Generates the list of lookup urls for the given url.  Returns an empty
        list if the given url is not valid.
    """
    if not self._IsValidUrl(url):
      logging.error('Invalid url %s' % url)
      return []
    return self._GetUrlCombinations(self._UrlCanonicalization(url))

  def _IsValidUrl(self, url):
    """Returns true if the url looks correct according to the Python url
       parser."""
    u = urlparse.urlsplit(url)
    return u.scheme != '' and u.hostname != ''
  
  def _UrlCanonicalization(self, url):
    """ Canonicalizes a url according to the safe browsing API specification
        PRE: Url must be valid
    """
    # TODO: handle international urls.
    # Repeatedly URL-unescape the URL until it has no more hex-encodings. 
    url = urllib.unquote(url)

    # To canonicalize the hostname: extract the hostname from the URL and do:
    parsed = urlparse.urlsplit(url)

    #  - Remove all leading and trailing dots .
    hostname = parsed.hostname.lstrip('.').rstrip('.')
    #  - Replace consecutive dots with a single dot.
    hostname = re.sub('\.+', '.', hostname)
    #  - If the hostname can be parsed as an IP address, it should be normalized
    #    to 4 dot-separated decimal values. The client should handle any legal
    #    IP address encoding, including octal, hex, and fewer than 4 components.
    # TODO
    #  - Lowercase the whole string.
    hostname = hostname.lower()

    # To canonicalize the path:
    #  - The sequences "/../" and "/./" in the path should be resolved, by
    #    replacing "/./" with "/", and removing "/../" along with the preceding 
    #    path component. 
    #  - Runs of consecutive slashes should be replaced with a single slash
    #    character. 
    path = parsed.path.replace('/./', '/')
    path = re.sub('\/+', '/', path)
    while path.find('/../') > 0:  # if path starts with /../ we don't remove it
      path = re.sub('/[^\/]*/../', '/', path)
  
    # The URL must include a path component, e.g. must have (at least) a
    # trailing slash
    if len(path) == 0:
      path = '/'  

    # Username, password and port are dropped at this point.  They are not
    # used in the current API.
    url = urlparse.urlunsplit((parsed.scheme, hostname, path, parsed.query,
                               parsed.fragment))

    # Precent-escape all characters in the URL which ASCII [32, 127] or %.
    # Escapes should use uppercase hex chars
    def PercentEscapeChar(c):
      if ord(c) <= 32 or ord(c) >= 127 or c == '%':
        return '%%%02X' % ord(c)
      return c
    return ''.join(map(PercentEscapeChar, url))

  def _GetUrlCombinations(self, url):
    """ Generates the required host suffix and path prefix combinations
        as specified by the safe browsing API.  This function may return an
        empty list.
        PRE: url is canonicalized.
    """ 
    parsed = urlparse.urlsplit(url)
    # First form all possible hostnames and paths and then return the cartisian
    # product of these two sets.

    # For the hostname, the client will try at most 5 different strings:
    #  - the exact hostname in the url 
    hostnames = [parsed.hostname]
    #  - up to 4 hostnames formed by starting with the last 5 components and
    #    successively removing the leading component. The top-level domain can
    #    be skipped.
    # E.g: a.b.c         => a.b.c, b.c
    #      a.b.c.d.e.f.g => a.b.c.d.e.f.g, c.d.e.f.g, d.e.f.g, e.f.g, f.g
    hnparts = parsed.hostname.split('.')[-5:]  # take the last five components
    if parsed.hostname.count('.') < 5:
      hnparts = hnparts[1:]
    while len(hnparts) > 1:
      hostnames.append('.'.join(hnparts))
      hnparts = hnparts[1:]  # remove the first element
  
    # For the path, the client will also try at most 6 different strings:
    #  - the exact path of the url, including query parameters
    paths = []
    if len(parsed.query) > 0:
      # TODO make sure that fragment identifiers e.g. #abc or not considered 
      paths.append('%s?%s' % (parsed.path, parsed.query))
    #  - the exact path without the query parameters
    paths.append(parsed.path)

    #  - the 4 paths formed by starting at the root (/) and successively
    #    appending path components, including a trailing slash. 
    paths.append('/')
    path_parts = parsed.path.rstrip('/').lstrip('/').split('/')[:4]
    if parsed.path.count('/') < 5:
      path_parts.pop()
    while path_parts:
      paths.append('/' + '/'.join(path_parts) + '/')
      path_parts.pop()  # remove last element
  
    # Return one url for each possible hostname and path
    urls = []
    for h in hostnames:
      for p in paths:
        urls.append(h + p)
    return urls
